import openai


def chat_ark(prompt, max_new_tokens=8192):
    openai.api_base = "https://ark.cn-beijing.volces.com/api/v3"
    openai.api_key = "93a67648-c2cd-4a51-99ba-c51114b537ee"
    messages = [
        {'role': 'user', 'content': prompt}
    ]
    response = openai.ChatCompletion.create(
        model="ep-m-20250221110554-mgvxq", # deepseek-r1-250120
        messages=messages,
        stream=False,
        max_tokens=max_new_tokens,
    )

    reasoning_content = response.choices[0].message.reasoning_content
    content = response.choices[0].message.content
    if not reasoning_content.startswith('<think>'):
        reasoning_content = '<think>\n' + reasoning_content
    if not reasoning_content.endswith('</think>'):
        reasoning_content += '\n</think>'
    content = content.strip()
    if not content.startswith('<answer>'):
        content = '<answer>\n' + content
    if not content.endswith('</answer>'):
        content += '\n</answer>'
    return f'{reasoning_content}\n{content}'